{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "ename": "TypeError",
     "evalue": "fsolve: there is a mismatch between the input and output shape of the 'func' argument 'objective_function'.Shape should be (2,) but it is (6,).",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[1;32m~\\AppData\\Local\\Temp\\ipykernel_27860\\2603983579.py\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[0;32m     58\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     59\u001b[0m \u001b[1;31m# 使用 fsolve 优化切换时间\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m---> 60\u001b[1;33m \u001b[0mt1\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mt2\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mfsolve\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mobjective_function\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m[\u001b[0m\u001b[0mt1_guess\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mt2_guess\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m     61\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     62\u001b[0m \u001b[1;31m# 内层循环求解 TFC\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32md:\\Conda\\lib\\site-packages\\scipy\\optimize\\_minpack_py.py\u001b[0m in \u001b[0;36mfsolve\u001b[1;34m(func, x0, args, fprime, full_output, col_deriv, xtol, maxfev, band, epsfcn, factor, diag)\u001b[0m\n\u001b[0;32m    158\u001b[0m                'diag': diag}\n\u001b[0;32m    159\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 160\u001b[1;33m     \u001b[0mres\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0m_root_hybr\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mfunc\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mx0\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0margs\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mjac\u001b[0m\u001b[1;33m=\u001b[0m\u001b[0mfprime\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0moptions\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    161\u001b[0m     \u001b[1;32mif\u001b[0m \u001b[0mfull_output\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    162\u001b[0m         \u001b[0mx\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mres\u001b[0m\u001b[1;33m[\u001b[0m\u001b[1;34m'x'\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32md:\\Conda\\lib\\site-packages\\scipy\\optimize\\_minpack_py.py\u001b[0m in \u001b[0;36m_root_hybr\u001b[1;34m(func, x0, args, jac, col_deriv, xtol, maxfev, band, eps, factor, diag, **unknown_options)\u001b[0m\n\u001b[0;32m    224\u001b[0m     \u001b[1;32mif\u001b[0m \u001b[1;32mnot\u001b[0m \u001b[0misinstance\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0margs\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mtuple\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    225\u001b[0m         \u001b[0margs\u001b[0m \u001b[1;33m=\u001b[0m \u001b[1;33m(\u001b[0m\u001b[0margs\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 226\u001b[1;33m     \u001b[0mshape\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mdtype\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0m_check_func\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;34m'fsolve'\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;34m'func'\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mfunc\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mx0\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0margs\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mn\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m(\u001b[0m\u001b[0mn\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    227\u001b[0m     \u001b[1;32mif\u001b[0m \u001b[0mepsfcn\u001b[0m \u001b[1;32mis\u001b[0m \u001b[1;32mNone\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    228\u001b[0m         \u001b[0mepsfcn\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mfinfo\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mdtype\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0meps\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32md:\\Conda\\lib\\site-packages\\scipy\\optimize\\_minpack_py.py\u001b[0m in \u001b[0;36m_check_func\u001b[1;34m(checker, argname, thefunc, x0, args, numinputs, output_shape)\u001b[0m\n\u001b[0;32m     36\u001b[0m                 \u001b[0mmsg\u001b[0m \u001b[1;33m+=\u001b[0m \u001b[1;34m\".\"\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     37\u001b[0m             \u001b[0mmsg\u001b[0m \u001b[1;33m+=\u001b[0m \u001b[1;34m'Shape should be %s but it is %s.'\u001b[0m \u001b[1;33m%\u001b[0m \u001b[1;33m(\u001b[0m\u001b[0moutput_shape\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mshape\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mres\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m---> 38\u001b[1;33m             \u001b[1;32mraise\u001b[0m \u001b[0mTypeError\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mmsg\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m     39\u001b[0m     \u001b[1;32mif\u001b[0m \u001b[0missubdtype\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mres\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mdtype\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0minexact\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     40\u001b[0m         \u001b[0mdt\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mres\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mdtype\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;31mTypeError\u001b[0m: fsolve: there is a mismatch between the input and output shape of the 'func' argument 'objective_function'.Shape should be (2,) but it is (6,)."
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "from scipy.optimize import fsolve\n",
    "from scipy.special import chebyt\n",
    "\n",
    "# 参数初始化\n",
    "ag = np.array([0, 0, -3.7114])  # 重力加速度 (m/s^2)\n",
    "Isp = 225  # 比冲 (s)\n",
    "g0 = 9.807  # 地球重力加速度 (m/s^2)\n",
    "Tmax = 13258.18  # 最大推力 (N)\n",
    "Tmin = 4971.81  # 最小推力 (N)\n",
    "alpha = 1 / (Isp * g0)  # 推力系数 (s/m)\n",
    "r0 = np.array([-200, 100, 1500])  # 初始位置 (m)\n",
    "v0 = np.array([85, 50, -65])  # 初始速度 (m/s)\n",
    "rf = np.array([0, 0, 0])  # 终止位置 (m)\n",
    "vf = np.array([0, 0, 0])  # 终止速度 (m/s)\n",
    "m0 = 1905  # 初始质量 (kg)\n",
    "tf_guess = 45  # 初始猜测的总时间 (s)\n",
    "\n",
    "# 定义时间区间\n",
    "t0 = 0\n",
    "tf = tf_guess\n",
    "N = 100  # 时间点数量\n",
    "t = np.linspace(t0, tf, N)\n",
    "\n",
    "# 初始化变量\n",
    "r = np.zeros((3, N))  # 位置\n",
    "v = np.zeros((3, N))  # 速度\n",
    "lambda_r = np.zeros((3, N))  # 位置伴随变量\n",
    "lambda_v = np.zeros((3, N))  # 速度伴随变量\n",
    "lambda_m = np.zeros(N)  # 质量伴随变量\n",
    "T = np.zeros(N)  # 推力\n",
    "\n",
    "# 初始条件\n",
    "r[:, 0] = r0\n",
    "v[:, 0] = v0\n",
    "\n",
    "# 使用 TFC 方法求解\n",
    "# 初始化 TFC 参数\n",
    "# 使用切比雪夫多项式\n",
    "m = 10  # 切比雪夫多项式的阶数\n",
    "z = np.cos(np.pi * np.arange(m + 1) / m)  # 切比雪夫节点\n",
    "h = np.array([chebyt(i)(z) for i in range(m + 1)]).T  # 切比雪夫多项式基函数\n",
    "\n",
    "# 定义切换时间 t1 和 t2\n",
    "t1_guess = 32.4  # 初始猜测的第一次切换时间\n",
    "t2_guess = 38.8  # 初始猜测的第二次切换时间\n",
    "\n",
    "# 外层循环优化切换时间\n",
    "def objective_function(x):\n",
    "    t1, t2 = x\n",
    "    # 在这里，你需要根据 t1 和 t2 的值，使用 TFC 方法求解整个过程\n",
    "    # 并计算最终位置和速度的误差作为目标函数的返回值\n",
    "    # 这里只是一个示例，需要根据论文中的具体公式实现\n",
    "    # 假设最终位置和速度的误差\n",
    "    final_position_error = rf - r[:, -1]\n",
    "    final_velocity_error = vf - v[:, -1]\n",
    "    return np.concatenate((final_position_error, final_velocity_error))\n",
    "\n",
    "# 使用 fsolve 优化切换时间\n",
    "t1, t2 = fsolve(objective_function, [t1_guess, t2_guess])\n",
    "\n",
    "# 内层循环求解 TFC\n",
    "for iter in range(100):  # 最大迭代次数\n",
    "    # 分段求解\n",
    "    for s in range(1, 4):\n",
    "        if s == 1:\n",
    "            ts = t0\n",
    "            te = t1\n",
    "        elif s == 2:\n",
    "            ts = t1\n",
    "            te = t2\n",
    "        else:\n",
    "            ts = t2\n",
    "            te = tf\n",
    "\n",
    "        # 根据当前段的切换逻辑设置推力 T\n",
    "        if s == 1:\n",
    "            T_segment = Tmax\n",
    "        elif s == 2:\n",
    "            T_segment = Tmin\n",
    "        else:\n",
    "            T_segment = Tmax\n",
    "\n",
    "        # 使用 TFC 方法求解当前段的 r, v, lambda_r, lambda_v, lambda_m\n",
    "        # 映射时间到切比雪夫节点\n",
    "        t_segment = np.linspace(ts, te, N)\n",
    "        z_segment = 2 * (t_segment - ts) / (te - ts) - 1\n",
    "\n",
    "        # 初始化自由函数的系数\n",
    "        xi = np.random.rand(m + 1, 3)  # 随机初始化\n",
    "\n",
    "        # 迭代求解\n",
    "        for k in range(100):  # 最大迭代次数\n",
    "            # 构造约束表达式\n",
    "            r_tfc = np.zeros((3, N))\n",
    "            v_tfc = np.zeros((3, N))\n",
    "            a_tfc = np.zeros((3, N))\n",
    "            for i in range(3):\n",
    "                r_tfc[i, :] = xi[:, i] @ h + boundary_conditions(z_segment, r0[i], rf[i], v0[i], vf[i])\n",
    "                v_tfc[i, :] = np.gradient(r_tfc[i, :], z_segment)\n",
    "                a_tfc[i, :] = np.gradient(v_tfc[i, :], z_segment)\n",
    "\n",
    "            # 计算残差\n",
    "            residuals = a_tfc - ag[:, None] - T_segment / m0 * lambda_v\n",
    "\n",
    "            # 更新自由函数的系数\n",
    "            J = np.zeros((3 * N, (m + 1) * 3))\n",
    "            for i in range(3):\n",
    "                J[i * N:(i + 1) * N, i * (m + 1):(i + 1) * (m + 1)] = np.gradient(np.gradient(h, z_segment), z_segment)\n",
    "            delta_xi = -np.linalg.pinv(J) @ residuals.flatten()\n",
    "            xi += delta_xi.reshape(m + 1, 3)\n",
    "\n",
    "            # 检查收敛性\n",
    "            if np.linalg.norm(delta_xi) < 1e-10:\n",
    "                break\n",
    "\n",
    "        # 更新位置和速度\n",
    "        r = r_tfc\n",
    "        v = v_tfc\n",
    "\n",
    "    # 检查收敛性\n",
    "    # 如果满足收敛条件，则退出循环\n",
    "    # 否则，更新自由函数的系数等参数，继续迭代\n",
    "\n",
    "# 输出结果\n",
    "print('最终切换时间:')\n",
    "print(f't1 = {t1}')\n",
    "print(f't2 = {t2}')\n",
    "print('最终时间:')\n",
    "print(tf)\n",
    "\n",
    "# 边界条件函数\n",
    "def boundary_conditions(z, r0, rf, v0, vf):\n",
    "    t0 = -1\n",
    "    tf = 1\n",
    "    t_star = z - t0\n",
    "    delta_t = tf - t0\n",
    "    bc = (1 + 2 * t_star**3 / delta_t**3 - 3 * t_star**2 / delta_t**2) * r0 + \\\n",
    "         (-2 * t_star**3 / delta_t**3 + 3 * t_star**2 / delta_t**2) * rf + \\\n",
    "         (t_star + t_star**2 / delta_t - 2 * t_star**3 / delta_t**2) * v0 + \\\n",
    "         (t_star**2 / delta_t - t_star**3 / delta_t**2) * vf\n",
    "    return bc"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
